import os
import re
from collections import defaultdict

class NaiveBayesClassifier:
    def __init__(self):
        self.classes = set()  # 类别集合
        self.class_counts = defaultdict(int)  # 每个类别的样本数
        self.word_counts = defaultdict(lambda: defaultdict(int))  # 类别-词-计数
        self.vocab = set()  # 词汇表
        self.total_docs = 0  # 总样本数

    def preprocess(self, text):
        """文本预处理：转小写、去特殊字符、分词"""
        text = text.lower()  # 小写化
        text = re.sub(r'[^a-zA-Z\s]', '', text)  # 保留字母和空格
        words = text.split()  # 分词
        return [word for word in words if len(word) > 2]  # 过滤短词

    def train(self, texts, labels):
        """训练模型：统计词频和类别数"""
        self.total_docs = len(texts)
        for text, label in zip(texts, labels):
            self.classes.add(label)
            self.class_counts[label] += 1  # 累加类别样本数
            words = self.preprocess(text)
            for word in words:
                self.vocab.add(word)
                self.word_counts[label][word] += 1  # 累加词频

    def predict(self, text):
        """预测文本类别"""
        words = self.preprocess(text)
        max_prob = -float('inf')
        best_class = None
        vocab_size = len(self.vocab)

        for label in self.classes:
            # 先验概率（拉普拉斯平滑）
            prior = (self.class_counts[label] + 1) / (self.total_docs + len(self.classes))
            # 条件概率乘积（拉普拉斯平滑）
            total_words = sum(self.word_counts[label].values())
            cond_prob = 1.0
            for word in words:
                count = self.word_counts[label].get(word, 0) + 1
                cond_prob *= count / (total_words + vocab_size)
            # 选择最大后验概率类别
            if prior * cond_prob > max_prob:
                max_prob = prior * cond_prob
                best_class = label
        return best_class

    def evaluate(self, test_texts, test_labels):
        """计算准确率"""
        correct = sum(1 for t, l in zip(test_texts, test_labels) if self.predict(t) == l)
        return correct / len(test_texts) if test_texts else 0.0